#ifndef SPATIALTF_PD_HPP
#define SPATIALTF_PD_HPP

#include "oneapi/dnnl/dnnl.h"

#include "c_types_map.hpp"
#include "primitive_desc.hpp"
#include "utils.hpp"

#define VDISPATCH_SPATIALTF(cond, msg, ...) \
VCONDCHECK(primitive, create, dispatch, spatialtf, (cond), \
status::unimplemented, "%s," msg, this->info(engine), \
##__VA_ARGS__)

#define VDISPATCH_SPATIALTF_SC(f, msg, ...) \
VCHECK(primitive, create, dispatch, spatialtf, (f), "%s," msg, \
this->info(engine), ##__VA_ARGS__)

namespace dnnl {
namespace impl {
struct spatialtf_gridgen_fwd_pd_t;
struct spatialtf_sampler_fwd_pd_t;

struct spatialtf_pd_t : public primitive_desc_t {
    static constexpr auto base_pkind = primitive_kind::spatialtf;
    const spatialtf_desc_t *desc() const { return &desc_; }
    const op_desc_t *op_desc() const override {
        return reinterpret_cast<const op_desc_t *>(this->desc());
    }

    status_t query(query_t what, int idx, void *result) const override {
        switch (what) {
            case query::prop_kind:
                *(prop_kind_t *)result = desc()->prop_kind;
            break;
            case query::primitive_kind:
                *(primitive_kind_t *)result = desc()->primitive_kind;
            break;
//            case query::alg_kind:
//                *(alg_kind_t *)result = desc()->alg_kind;
            break;
            default: return primitive_desc_t::query(what, idx, result);
        }
        return status::success;
    }

    const int src0_ndims() const { return desc_.src0_desc.ndims; }
    const void get_src0_dims(int *dims) const {
        for (int i = 0; i < src0_ndims(); i++) {
            dims[i] = desc_.src0_desc.dims[i];
        }
    }

    const memory_desc_t *src_md(
        int index = 0, bool user_input = false) const override {
        if (index == 0) return user_input ? &desc()->src1_desc : &src1_md_;
        return &glob_zero_md;
    }
    const memory_desc_t *dst_md(
            int index = 0, bool user_input = false) const override {
        if (index == 0) return user_input ? &desc()->dst_desc : &dst_md_;
        return &glob_zero_md;
    }



protected:
    memory_desc_t src1_md_;
    spatialtf_desc_t desc_;
    const spatialtf_gridgen_fwd_pd_t *hint_fwd_pd_grid;
    const spatialtf_sampler_fwd_pd_t *hint_fwd_pd_sampler;

    memory_desc_t dst_md_;

    spatialtf_pd_t(const spatialtf_desc_t *adesc, const primitive_attr_t *attr,
            const spatialtf_gridgen_fwd_pd_t *hint_fwd_pd)
        : primitive_desc_t(attr, base_pkind)
        , desc_(*adesc)
        , hint_fwd_pd_grid(hint_fwd_pd)
        , dst_md_(desc_.dst_desc)
        {
        src1_md_ = desc_.src1_desc;
        }

private:
    const memory_desc_t &dst_desc() const { return dst_md_; }

    };

struct spatialtf_gridgen_fwd_pd_t : public spatialtf_pd_t {
    typedef spatialtf_gridgen_fwd_pd_t base_class;
    typedef spatialtf_gridgen_fwd_pd_t hint_class;

    arg_usage_t arg_usage(int arg) const override {
        if (arg == DNNL_ARG_SRC_1) return arg_usage_t::input;

        if (arg == DNNL_ARG_DST) return arg_usage_t::output;

        if (arg == DNNL_ARG_WORKSPACE && (!types::is_zero_md(workspace_md())))
            return arg_usage_t::output;

        return primitive_desc_t::arg_usage(arg);
    }

    int n_inputs() const override { return 1 ; }
    int n_outputs() const override {
        return 1;
    }

protected:
    spatialtf_gridgen_fwd_pd_t(const spatialtf_desc_t *adesc, const primitive_attr_t *attr,
            const spatialtf_gridgen_fwd_pd_t *hint_fwd_pd):
        spatialtf_pd_t(adesc, attr, hint_fwd_pd){
        }
};

struct spatialtf_gridgen_bwd_pd_t : public spatialtf_pd_t {};

struct spatialtf_sampler_fwd_pd_t : public spatialtf_pd_t {};

struct spatialtf_sampler_bwd_pd_t : public spatialtf_pd_t {};


} // namespace impl
} // namespace dnnl

#endif //SPATIALTF_PD_HPP
